package pl.streamsoft.test.helpers;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Properties;

public class DatabaseHelper {
	static final String DATABASE_URL = "jdbc:postgresql://localhost/test";
	static final String DATABASE_USER = "postgres";
	static final String DATABASE_PASSWORD = "postgres";

	private Connection connection;

	protected void connect() throws SQLException {
		Properties connectionProperties = new Properties();
		connectionProperties.setProperty("user", DATABASE_USER);
		connectionProperties.setProperty("password", DATABASE_PASSWORD);

		connection = DriverManager.getConnection(DATABASE_URL,
				connectionProperties);
	}

	public void disconnect() throws SQLException {
		connection.close();
	}

	public void truncateTable(String tableName) {
		try {
			if (connection == null || connection.isClosed()) {
				connect();
			}
			Statement statement = connection.createStatement();

			statement.executeUpdate("TRUNCATE " + tableName + " CASCADE");
		} catch (SQLException e) {
			e.printStackTrace();
		}
	}

	public int getRowsCount(String tableName) {
		try {
			if (connection == null || connection.isClosed()) {
				connect();
			}
			Statement statement = connection.createStatement();

			ResultSet resultSet = statement
					.executeQuery("SELECT count(*) FROM " + tableName);
			if (resultSet != null && resultSet.next()) {
				return resultSet.getInt(1);
			}
		} catch (SQLException e) {
			e.printStackTrace();
		}
		return 0;
	}
}
